package com.boge.ai.utils;

import com.boge.ai.entity.MilvusEntity;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MilvusExt;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import okhttp3.*;
import org.apache.commons.math3.linear.RealVector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Milvus 向量数据库的公共方法
 */
public class MilvusLLMUtils {


    private MilvusServiceClient client ;

    public MilvusLLMUtils() {
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();
        client = new MilvusServiceClient(connectParam);

    }


    /**
     * 创建集合
     * @throws Exception
     */
    public  void createCollection() throws Exception {

        List<FieldType> fieldTypes = Arrays.asList(
                FieldType.newBuilder()
                        .withName(MilvusEntity.Field.ID)
                        .withDescription("主键ID")
                        .withDataType(DataType.Int64)
                        .withPrimaryKey(true)
                        .withAutoID(true)
                        .build(),
                FieldType.newBuilder()
                        .withName(MilvusEntity.Field.FEATURE)
                        .withDescription("特征向量")
                        .withDataType(DataType.FloatVector)
                        .withDimension(MilvusEntity.FEATURE_DIM) // 设置向量维度
                        .build(),
                FieldType.newBuilder()
                        .withName(MilvusEntity.Field.TEXT)
                        .withDescription("输入数据")
                        .withDataType(DataType.VarChar)
                        .withTypeParams(Collections.singletonMap("max_length", "65535"))
                        .build(),
                FieldType.newBuilder()
                        .withName(MilvusEntity.Field.OUTPUT)
                        .withDescription("问题答案数据")
                        .withDataType(DataType.VarChar)
                        .withTypeParams(Collections.singletonMap("max_length", "65535"))
                        .build());

        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withDescription("rag collection")
                .withShardsNum(MilvusEntity.SHARDS_NUM)
                .withFieldTypes(fieldTypes)
                .build();
        client.createCollection(createCollectionReq);
        // 同时给向量创建对应的索引
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withFieldName(MilvusEntity.Field.FEATURE) // 向量字段名
                .withIndexType(IndexType.IVF_FLAT) // 使用IVF_FLAT索引类型
                .withMetricType(MetricType.L2) // 指定度量类型，如L2距离
                .withExtraParam("{\"nlist\":128}") // 根据索引类型提供额外参数，比如nlist
                .build();

        client.createIndex(createIndexParam);
    }

    /**
     * 插入数据到向量数据库
     * @throws Exception
     */
    public  void insertVectoryData(List<Float> vectorParam,String text,String output) throws Exception {

        createCollection();
        List<List<Float>> floats = new ArrayList<>();
        floats.add(vectorParam);
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field(MilvusEntity.Field.FEATURE, floats));
        fields.add(new InsertParam.Field(MilvusEntity.Field.TEXT, Arrays.asList(text)));
        fields.add(new InsertParam.Field(MilvusEntity.Field.OUTPUT, Arrays.asList(output)));


        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withFields(fields)
                .build();
        client.insert(insertParam);
    }


    /**
     * 根据向量检索信息
     * @param searchVectors
     * @return
     * @throws Exception
     */
    public  SearchResultsWrapper search(List<Float> searchVectors) throws Exception {
        List<List<Float>> floats = new ArrayList<>();
        floats.add(searchVectors);

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withMetricType(MetricType.L2)// 使用 L2 距离作为相似度度量
                .withTopK(3) // 返回最接近的前3个结果
                .withVectors(floats)
                .withVectorFieldName(MilvusEntity.Field.FEATURE) // 向量字段名
                .withOutFields(Arrays.asList(MilvusEntity.Field.ID,MilvusEntity.Field.OUTPUT,MilvusEntity.Field.TEXT)) // 需要返回的字段
                .build();
        SearchResults data = client.search(searchParam).getData();
        if(data != null) {
            SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());

            return resultsWrapper;
        }
        return null;
    }

}
